#include "xyz2rgb.h"

XYZ2RGB::XYZ2RGB(std::string engine_path, float conf_thres, float nms_thresh, ros::NodeHandle n)
: DetectNode(engine_path, conf_thres, nms_thresh), n_(n)
{   
    getParam();

    cloud_.reset(new pcl::PointCloud<PointT>());
    cloud_voxel_.reset(new pcl::PointCloud<PointT>());
    cloud_rgb_.reset(new pcl::PointCloud<pcl::PointXYZRGB>());
    
    msg_cloud_new_.reset(new sensor_msgs::PointCloud2());

    pub_img_ = n_.advertise<sensor_msgs::Image>(topic_img_new_, 10);
    pub_cloud_ = n_.advertise<sensor_msgs::PointCloud2>(topic_cloud_new_, 10);
    
    // 5 同步函数定义
    sub_cloud_.subscribe(n_, topic_cloud_, 10);
    sub_img_.subscribe(n_, topic_img_, 10);

    sync_ = std::make_shared<Sync>( MySyncPolicy(10), sub_cloud_, sub_img_ );
    sync_->registerCallback( boost::bind(&XYZ2RGB::callBack, this, _1, _2) ); 
}

void XYZ2RGB::callBack(const sensor_msgs::PointCloud2ConstPtr& msg_cloud, const sensor_msgs::ImageConstPtr& msg_img)
{      
    ROS_WARN("-------\n");
    img_ = cv_bridge::toCvShare(msg_img, "bgr8")->image;
    cv::Mat img_det = img_.clone();
    std::vector<Detection> res;
    
    inference(img_det, res);
    draw_bbox(img_det, res);

    if(is_filter_)
    {
        pcl::fromROSMsg(*msg_cloud, *cloud_);  // 雷达坐标系,x为深度
        filterVoxel(cloud_, cloud_voxel_, filter_size_);
    }
    else
    {
        pcl::fromROSMsg(*msg_cloud, *cloud_voxel_);
    }
    
    cloud_rgb_->clear();
    
    // 点云库坐标变换
    auto t1 = std::chrono::high_resolution_clock::now();
    cloudProcess(cloud_voxel_, cloud_rgb_);
    auto t2 = std::chrono::high_resolution_clock::now();
    auto cost_time = std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1);
    // std::cout << "cost_time: " << cost_time.count() << std::endl;

    // 遍历点进行坐标变换,上面速度是下面的9-10倍
    // auto t1 = std::chrono::high_resolution_clock::now();
    // pointProcess(cloud_voxel_, cloud_rgb_);
    // auto t2 = std::chrono::high_resolution_clock::now();
    // auto duration2 = std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1);

    pcl::toROSMsg(*cloud_rgb_, *msg_cloud_new_);
    msg_cloud_new_->header.stamp = msg_cloud->header.stamp;
    msg_cloud_new_->header.frame_id = frame_id_;
    pub_cloud_.publish(msg_cloud_new_);
    
    msg_img_new_ = cv_bridge::CvImage(msg_cloud->header, "bgr8", img_det).toImageMsg();
    pub_img_.publish(msg_img_new_);

    ROS_INFO("Detection and ranging is %.2f fps ...", 1000.0 / (cost_time.count() + cost_time_));
}


void XYZ2RGB::filterVoxel(const pcl::PointCloud<PointT>::Ptr& cloud_in, pcl::PointCloud<PointT>::Ptr& cloud_out, float size)
{
    pcl::VoxelGrid<PointT> sor; // 创建体素滤波器
    sor.setInputCloud(cloud_in);          // 输入点云
    // sor.setLeafSize(0.01f, 0.01f, 0.01f); //设置滤波时创建的体素体积为1cm的立方体,m
    sor.setLeafSize(size, size, size);
    sor.filter(*cloud_out);
}

void XYZ2RGB::Rt2T(Eigen::Isometry3d& T)
{
    // R,t->Tcl 
    Eigen::Vector3d R_euler(roll_, pitch_, yaw_); // 旋转
    R_euler = R_euler / 180 * M_PI;
    
    Eigen::Vector3d t(x_, y_, z_);  // 平移
    Eigen::Matrix3d R;
    R = Eigen::AngleAxisd(R_euler[0], Eigen::Vector3d::UnitZ()) * 
        Eigen::AngleAxisd(R_euler[1], Eigen::Vector3d::UnitY()) * 
        Eigen::AngleAxisd(R_euler[2], Eigen::Vector3d::UnitX());
    T.rotate(R);
    T.pretranslate (t);
}

void XYZ2RGB::cloudProcess(const pcl::PointCloud<PointT>::Ptr& cloud_in, pcl::PointCloud<pcl::PointXYZRGB>::Ptr& cloud_rgb)
{
    pcl::PointCloud<PointT>::Ptr cloud_cam(new pcl::PointCloud<PointT>);
    pcl::PointCloud<PointT>::Ptr cloud_uv(new pcl::PointCloud<PointT>);

    pcl::transformPointCloud(*cloud_in, *cloud_cam, Eigen::Affine3d(Tcl_)); // 点云转到相机坐标下
    pcl::transformPointCloud(*cloud_cam, *cloud_uv, Eigen::Affine3d(K_));   // 像素坐标但未归一化
    Eigen::Vector3d p_uv(0,0,1);
    pcl::PointXYZRGB p_rgb;
    for (int i = 0; i < cloud_cam->points.size(); ++i)
    {
        if(cloud_cam->points[i].z < 0)   // 激光下的x变到相机下的z,这里要判断z是否小于0
            continue;
        
        p_rgb.x = cloud_cam->points[i].x;
        p_rgb.y = cloud_cam->points[i].y;
        p_rgb.z = cloud_cam->points[i].z;

        p_rgb.b = 255;
        p_rgb.g = 0;
        p_rgb.r = 0;
        
        p_uv[0] = cloud_uv->points[i].x / cloud_uv->points[i].z;
        p_uv[1] = cloud_uv->points[i].y / cloud_uv->points[i].z;

        if(p_uv[0] >= 0 && p_uv[0] <= img_.cols &&
        p_uv[1] >= 0 && p_uv[1] <= img_.rows)
        {
            // p_rgb.b = img_.at<cv::Vec3b>(int(p_uv[1]), int(p_uv[0]))[0];
            // p_rgb.g = img_.at<cv::Vec3b>(int(p_uv[1]), int(p_uv[0]))[1];
            // p_rgb.r = img_.at<cv::Vec3b>(int(p_uv[1]), int(p_uv[0]))[2];
            
            p_rgb.b = img_.ptr<uchar>(int(p_uv[1]))[int(p_uv[0]) * 3];
            p_rgb.g = img_.ptr<uchar>(int(p_uv[1]))[int(p_uv[0]) * 3 + 1];
            p_rgb.r = img_.ptr<uchar>(int(p_uv[1]))[int(p_uv[0]) * 3 + 2];

        }
        cloud_rgb->points.emplace_back(p_rgb);
    }

}


void XYZ2RGB::pointProcess(const pcl::PointCloud<PointT>::Ptr& cloud_in, pcl::PointCloud<pcl::PointXYZRGB>::Ptr& cloud_rgb)
{
    Eigen::Vector3d p_uv(0,0,1), p_cam(0,0,0), p_lidar(0,0,0);
    pcl::PointXYZRGB p_rgb;
    for (auto it : cloud_in->points)
    {
        
        p_lidar[0] = it.x;
        p_lidar[1] = it.y;
        p_lidar[2] = it.z;

        p_cam = Tcl_ * p_lidar;
        if(p_cam[2] < 0)
            continue;

        p_uv = K_ * p_cam / p_cam[2]; 
        p_rgb.x = p_cam[0];
        p_rgb.y = p_cam[1];
        p_rgb.z = p_cam[2];

        p_rgb.b = 255;
        p_rgb.g = 0;
        p_rgb.r = 0;

        if(p_uv[0] >= 0 && p_uv[0] <= img_.cols &&
        p_uv[1] >= 0 && p_uv[1] <= img_.rows)
        {
            p_rgb.b = img_.at<cv::Vec3b>(p_uv[1], p_uv[0])[0];
            p_rgb.g = img_.at<cv::Vec3b>(p_uv[1], p_uv[0])[1];
            p_rgb.r = img_.at<cv::Vec3b>(p_uv[1], p_uv[0])[2];
        }
        
        cloud_rgb->points.emplace_back(p_rgb);
    }
}


void XYZ2RGB::getParam()
{
    // detector_->weight_path_ = detector_->pkg_path_ + "/weights/n_cpu.torchscript";
    n_.getParam("filter_size", filter_size_);
    n_.param<float>("filter_size", filter_size_, 0.2);
    
    n_.getParam("is_filter", is_filter_);
    n_.param<bool>("is_filter", is_filter_, false);


    n_.getParam("frame_id", frame_id_);
    n_.param<std::string>("frame_id", frame_id_, "map"); // 左右二维码中心距
    
    n_.getParam("topic_img", topic_img_);
	n_.param<std::string>("topic_img", topic_img_, "/image_left");

	n_.getParam("topic_img_new", topic_img_new_);
	n_.param<std::string>("topic_img_new", topic_img_new_, "/detect/image");

    n_.getParam("topic_cloud", topic_cloud_);
    n_.param<std::string>("topic_cloud", topic_cloud_, "/velodyne_points");
    
    n_.getParam("topic_cloud_new", topic_cloud_new_);
    n_.param<std::string>("topic_cloud_new", topic_cloud_new_, "/detect/rgb_points");



    // ------------内参---------------------
    n_.getParam("fx", fx_);
    n_.param<float>("fx", fx_, 718.856); 
    n_.getParam("fy", fy_);
    n_.param<float>("fy", fy_, 718.856); 
    n_.getParam("cx", cx_);
    n_.param<float>("cx", cx_, 607.1928); 
    n_.getParam("cy", cy_);
    n_.param<float>("cy", cy_, 185.2157); 
    
    n_.getParam("k1", k1_);
    n_.param<float>("k1", k1_, 0.0); 
    n_.getParam("k2", k2_);
    n_.param<float>("k2", k2_, 0.0); 
    n_.getParam("k3", k3_);
    n_.param<float>("k3", k3_, 0.0); 
    
    n_.getParam("p1", p1_);
    n_.param<float>("p1", p1_, 0.0); 
    n_.getParam("p2", p2_);
    n_.param<float>("p2", p2_, 0.0); 

    //  --------Tcl--------
    n_.getParam("roll", roll_);
    n_.param<float>("roll", roll_, 90.); 
    n_.getParam("pitch_", pitch_);
    n_.param<float>("pitch", pitch_, -90.); 
    n_.getParam("yaw", yaw_);
    n_.param<float>("yaw", yaw_, 0.0); 
    
    n_.getParam("x", x_);
    n_.getParam("y", y_);
    n_.getParam("z", z_);

    n_.param<float>("x", x_, -1.198459927713e-02);  
    n_.param<float>("y", y_, 5.403984729748e-02); 
    n_.param<float>("z", z_, -2.921968648686e-01); 

    // n_.param<float>("x", x_, 0); 
    // n_.param<float>("y", y_, 0);
    // n_.param<float>("z", z_, 0); 

    cameraMatrix_ = (cv::Mat_<float>(3, 3) << fx_, 0., cx_,
                                            0, fy_, cy_,
                                            0, 0., 1);
    
    distCoeffs_ = (cv::Mat_<float>(5, 1) << k1_, k2_, p1_, p2_, k3_);  // 初始化加（）
    
    Eigen::Matrix3d K;
    cv::cv2eigen(cameraMatrix_, K);
    K_.rotate(K);
    K_.pretranslate(Eigen::Vector3d (0, 0, 0));

    Rt2T(Tcl_);

    std::cout << "\n---------------\n" << std::endl;
    std::cout << " --cameraMatrix: \n" << K_.matrix() << std::endl;
    std::cout << "\n --distCoeffs: \n"<< distCoeffs_.at<float>(0,0) << " " <<
                                    distCoeffs_.at<float>(1,0) << " " <<
                                    distCoeffs_.at<float>(2,0) << " " <<
                                    distCoeffs_.at<float>(3,0) << " " <<
                                    distCoeffs_.at<float>(4,0) << " " << std::endl;
    std::cout << "\n --T_lidar2cam:\n" << Tcl_.matrix() << std::endl;
    
    std::cout << "\n-------------------\n" << std::endl;
    
    std::cout << " --filter_size : "<< filter_size_ << std::endl;
    std::cout << " --is_filter: " << is_filter_ << std::endl; 
    // std::cout << " --is_draw_distance : "<< is_draw_distance_ << std::endl;
    // std::cout << " --is_xyz2rgb: " << is_xyz2rgb_ << std::endl; 
    // std::cout << " --find_point_radius : "<< find_point_radius_ << std::endl;

    std::cout << " --topic_img_: " << topic_img_ << std::endl; 
    std::cout << " --topic_cloud: " << topic_cloud_ << std::endl; 
    
    std::cout << " --frame_id: " << frame_id_ << std::endl; 
    std::cout << " --topic_img_new: " << topic_img_new_ << std::endl; 
    std::cout << " --topic_cloud_new: " << topic_cloud_new_ << std::endl; 
    // std::cout << " --topic_box: " << topic_box_ << std::endl; 

    std::cout << "\n---------------\n" << std::endl;
}